# Copyright (c) 2025 Songwufeng
# [Software Name] is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import numpy as np
import cv2
from ultralytics import YOLO
import win32gui
import win32con
import time
import threading
import mss
import os
from pynput import keyboard
import ctypes
import pyautogui
import torch
# PID 控制器类
class PID:
    def __init__(self, P=0.2, I=0.01, D=0.1):
        self.kp, self.ki, self.kd = P, I, D
        self.uPrevious, self.uCurent = 0, 0
        self.setValue, self.lastErr, self.errSum = 0, 0, 0
        self.errSumLimit = 10

    def pidPosition(self, setValue, curValue):
        err = setValue - curValue
        dErr = err - self.lastErr
        self.errSum += err
        outPID = self.kp * err + self.ki * self.errSum + self.kd * dErr
        self.lastErr = err
        return outPID

# 罗技驱动类
class LOGITECH:
    def __init__(self):
        self.dll = None
        self.state = False
        self.load_dll()

    def load_dll(self):
        try:
            dll_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logitech.driver.dll')
            self.dll = ctypes.CDLL(dll_path)
            self.dll.device_open.restype = ctypes.c_int
            result = self.dll.device_open()
            self.state = (result == 1)
            self.dll.moveR.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_bool]
            self.dll.moveR.restype = None
        except FileNotFoundError:
            print(f'错误, 找不到 DLL 文件')
        except OSError as e:
            print(f'错误, 加载 DLL 文件失败: {e}')

    def mouse_move(self, end_xy, min_xy=2):
        if not self.state:
            return

        end_x, end_y = end_xy
        pid_x = PID()
        pid_y = PID()

        new_x, new_y = pyautogui.position()
        move_x = pid_x.pidPosition(end_x, new_x)
        move_y = pid_y.pidPosition(end_y, new_y)

        move_x = np.clip(move_x, -min_xy if move_x < 0 else min_xy, None).astype(int)
        move_y = np.clip(move_y, -min_xy if move_y < 0 else min_xy, None).astype(int)

        self.dll.moveR(move_x, move_y, True)

# 平滑坐标的函数，使用简单移动平均
class CoordinateSmoother:
    def __init__(self, window_size=5):
        self.window_size = window_size
        self.x_buffer = []
        self.y_buffer = []

    def smooth_coordinate(self, x, y):
        self.x_buffer.append(x)
        self.y_buffer.append(y)

        if len(self.x_buffer) > self.window_size:
            self.x_buffer.pop(0)
            self.y_buffer.pop(0)

        smoothed_x = int(np.mean(self.x_buffer))
        smoothed_y = int(np.mean(self.y_buffer))

        return smoothed_x, smoothed_y

# 全局变量
target_center_coordinates = []
latest_screenshot = None
lock = threading.Lock()
logitech = LOGITECH()
coordinate_smoother = CoordinateSmoother()
shift_pressed = False

# 初始化模型和窗口
def init_model_and_window(model_path, window_name, width, height, shift_x):
    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), model_path)
    if not os.path.exists(model_path):
        print(f"模型文件 {model_path} 不存在，请检查。")
        return None, None

    model = YOLO(model_path)
    model.fuse()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model.to(device)

    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(window_name, width, height)
    time.sleep(0.1)
    hwnd = win32gui.FindWindow(None, window_name)
    win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)
    return model, hwnd

# 截图线程函数
def screenshot_thread(width, height, shift_x, screenshot_interval):
    global latest_screenshot
    last_screenshot_time = time.time()
    with mss.mss() as sct:
        monitor = sct.monitors[0]
        start_x = (monitor["width"] - width) // 2 + shift_x
        start_y = (monitor["height"] - height) // 2
        monitor_area = {
            "left": start_x,
            "top": start_y,
            "width": width,
            "height": height
        }
        while True:
            current_time = time.time()
            if current_time - last_screenshot_time > screenshot_interval:
                try:
                    sct_img = sct.grab(monitor_area)
                    screenshot_np = np.array(sct_img)
                    screenshot_np = cv2.cvtColor(screenshot_np, cv2.COLOR_BGRA2BGR)
                    with lock:
                        latest_screenshot = screenshot_np
                    last_screenshot_time = current_time
                except Exception as e:
                    print(f"截图时发生错误: {e}")
            time.sleep(0.001)

# 实时存放目标中心坐标点的函数
def store_target_center_coordinates(x_center, y_center):
    global target_center_coordinates
    target_center_coordinates.append((x_center, y_center))
    if len(target_center_coordinates) > 100:
        target_center_coordinates.pop(0)

# 目标检测和处理函数
def detect_and_process(model, width, height, shift_x, window_name):
    history_boxes = []
    history_confidences = []
    history_size = 5
    topmost_interval = 1
    last_topmost_time = time.time()

    with mss.mss() as sct:
        monitor = sct.monitors[0]
        start_x = (monitor["width"] - width) // 2 + shift_x
        start_y = (monitor["height"] - height) // 2

    while True:
        current_time = time.time()
        if current_time - last_topmost_time > topmost_interval:
            hwnd = win32gui.FindWindow(None, window_name)
            if hwnd != 0:
                win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0, win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)
            last_topmost_time = current_time

        with lock:
            if latest_screenshot is None:
                continue
            screenshot_np = latest_screenshot.copy()

        try:
            results = model(screenshot_np, verbose=False, half=True)
        except Exception as e:
            print(f"目标检测时发生错误: {e}")
            continue

        max_conf = 0
        best_box = None
        drawn_y_ranges = []

        for result in results:
            boxes = result.boxes.cpu().numpy()
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0].astype(int)
                is_overlap = any(y1 < y2_ and y2 > y1_ for y1_, y2_ in drawn_y_ranges)
                if is_overlap:
                    continue

                conf = box.conf[0]
                history_confidences.append(conf)
                if len(history_confidences) > history_size:
                    history_confidences.pop(0)
                smoothed_conf = np.mean(history_confidences)

                if smoothed_conf > max_conf:
                    max_conf = smoothed_conf
                    best_box = box
                drawn_y_ranges.append((y1, y2))

        result_image = screenshot_np.copy()

        if best_box is not None:
            x1, y1, x2, y2 = best_box.xyxy[0].astype(int)
            current_box = np.array([x1, y1, x2, y2])
            history_boxes.append(current_box)
            if len(history_boxes) > history_size:
                history_boxes.pop(0)
            avg_box = np.mean(history_boxes, axis=0).astype(int)
            x1, y1, x2, y2 = avg_box

            cls = int(best_box.cls[0])
            class_name = results[0].names[cls]
            conf = max_conf

            cv2.rectangle(result_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f'{class_name}: {conf:.2f}'
            cv2.putText(result_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

            x_center = (x1 + x2) // 2
            y_center = (y1 + y2) // 2

            screen_x_center = start_x + x_center
            screen_y_center = start_y + y_center

            store_target_center_coordinates(screen_x_center, screen_y_center)

            if shift_pressed:
                smoothed_x, smoothed_y = coordinate_smoother.smooth_coordinate(screen_x_center, screen_y_center)
                logitech.mouse_move((smoothed_x, smoothed_y))

        cv2.imshow(window_name, result_image)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

# 封装获取目标中心坐标的函数
def get_target_center_coordinates():
    global target_center_coordinates
    return target_center_coordinates

# 键盘事件处理函数
def on_press(key):
    global shift_pressed
    if key == keyboard.Key.shift:
        shift_pressed = True

def on_release(key):
    global shift_pressed
    if key == keyboard.Key.shift:
        shift_pressed = False
    if key == keyboard.Key.esc:
        return False

def main():
    model_path = 'best.pt'
    window_name = 'Detection Result'
    width, height = 320, 320
    shift_x = 120
    screenshot_interval = 0.1

    import torch  # 导入 torch 库
    model, hwnd = init_model_and_window(model_path, window_name, width, height, shift_x)
    if model is None:
        return

    screenshot_t = threading.Thread(target=screenshot_thread, args=(width, height, shift_x, screenshot_interval))
    screenshot_t.daemon = True
    screenshot_t.start()

    keyboard_listener = keyboard.Listener(on_press=on_press, on_release=on_release)
    keyboard_listener.start()

    try:
        detect_and_process(model, width, height, shift_x, window_name)
    except KeyboardInterrupt:
        print("程序被用户手动中断。")
    except Exception as e:
        print(f"程序运行过程中发生错误: {e}")
    finally:
        cv2.destroyAllWindows()
        keyboard_listener.stop()
        print("存储的目标中心坐标点:", get_target_center_coordinates())

if __name__ == "__main__":
    main()